import torch
import numpy as np
from termcolor import colored
import pdb
class NumPrint:
    def __init__(self):
        self.print_enable = True
    def set_print_enable(self, enable):
        self.print_enable = enable
    def print(self, tag, origarr, option="sl",l_dim=-1, force_full_show=False, round_digits=3, override_print=False,todivice=None):
        try:
            if origarr is None:
                if self.print_enable:
                    print(tag,"is",origarr)
                return
            if not self.print_enable and not override_print:
                return
            arr = origarr
            
            if isinstance(origarr, torch.Tensor):
                arr = origarr.detach().cpu().numpy()
            output = f"{colored(tag,'green',attrs=['bold'])}"
            if not hasattr(arr, 'shape') or not hasattr(arr, 'ndim'):
                output += f" {arr}"
                print(output)
                return
            if "s" in option:
                output += f"{colored(' shape', 'cyan',attrs=['bold'])}{arr.shape}"
                
            if "l" in option or "c" in option:
                if "l" in option:
                    selection = [0] * arr.ndim
                    selection[l_dim] = slice(None)
                    arr_to_show = arr[tuple(selection)]
                else:
                    arr_to_show = arr
                if arr_to_show.flatten().shape[0] > 16 and not force_full_show:
                    print(colored(f"{tag}","green",attrs=['bold']),colored(f"toshowshape {arr_to_show.shape} origshape {arr.shape} first {arr_to_show.flatten()[0]}","yellow"))
                    return
                elif arr_to_show.dtype in [bool,int]:
                    output += f" {arr_to_show.astype(int)}"
                    all0 = np.all(arr == 0)
                    all1 = np.all(arr == 1)
                    if all0:
                        output += f" [all0:{all0}]"
                    elif all1:
                        output += f" [all1:{all1}]"
                else:
                    rounded_arr = np.round(arr_to_show, round_digits)
                    list_arr = list(rounded_arr)
                    output += f" {list_arr}"
            if hasattr(arr, 'device'):
                output += f" {arr.device}"
            output += f" {str(type(arr))}"
            print(output)
            if todivice is not None and isinstance(origarr, torch.Tensor):
                # print(f"origarr.device {origarr.device} todivice {todivice}")
                return origarr.to(todivice)
            return origarr
        except Exception as e:
            print(e)
            return origarr
p=NumPrint()